﻿#!/usr/bin/env python3

import smbus2,time,math
import ctypes



def set_x_accel_offset(a_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x06,ctypes.c_int8(a_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x07,ctypes.c_int8(a_offset).value)

def set_y_accel_offset(a_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x08,ctypes.c_int8(a_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x09,ctypes.c_int8(a_offset).value)

def set_z_accel_offset(a_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x0A,ctypes.c_int8(a_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x0B,ctypes.c_int8(a_offset).value)

def set_x_gyro_offset(g_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x13,ctypes.c_int8(g_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x14,ctypes.c_int8(g_offset).value)

def set_y_gyro_offset(g_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x15,ctypes.c_int8(g_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x16,ctypes.c_int8(g_offset).value)

def set_z_gyro_offset(g_offset):
    bus.write_byte_data(MPU6050_ADDR, 0x17,ctypes.c_int8(g_offset >> 8).value)
    bus.write_byte_data(MPU6050_ADDR, 0x18,ctypes.c_int8(g_offset).value)


x_accel_offset = int(-1588)
y_accel_offset = int(1520.875)
z_accel_offset = int(-755.5)
x_gyro_offset = int(16.5)
y_gyro_offset = int(-13.25)
z_gyro_offset = int(-25)



def MPU6050_start():
    # alter sample rate (stability)
    samp_rate_div = 0 # sample rate = 8 kHz/(1+samp_rate_div)
    bus.write_byte_data(MPU6050_ADDR, SMPLRT_DIV, samp_rate_div)
    time.sleep(0.1)
    # reset all sensors
    bus.write_byte_data(MPU6050_ADDR,PWR_MGMT_1,0x00)
    time.sleep(0.1)
    # power management and crystal settings
    bus.write_byte_data(MPU6050_ADDR, PWR_MGMT_1, 0x01)
    time.sleep(0.1)
    #Write to Configuration register
    bus.write_byte_data(MPU6050_ADDR, CONFIG, 0)
    time.sleep(0.1)
    #Write to Gyro configuration register
    gyro_config_sel = [0b00000,0b010000,0b10000,0b11000] # byte registers
    gyro_config_vals = [250.0,500.0,1000.0,2000.0] # degrees/sec
    gyro_indx = 0
    bus.write_byte_data(MPU6050_ADDR, GYRO_CONFIG, int(gyro_config_sel[gyro_indx]))
    time.sleep(0.1)
    #Write to Accel configuration register
    accel_config_sel = [0b00000,0b01000,0b10000,0b11000] # byte registers
    accel_config_vals = [2.0,4.0,8.0,16.0] # g (g = 9.81 m/s^2)
    accel_indx = 0                            
    bus.write_byte_data(MPU6050_ADDR, ACCEL_CONFIG, int(accel_config_sel[accel_indx]))
    time.sleep(0.1)
    # interrupt register (related to overflow of data [FIFO])
    bus.write_byte_data(MPU6050_ADDR, INT_ENABLE, 1)
    time.sleep(0.1)

    # x_accel_offset = int(-1588)
    # y_accel_offset = int(1520.875)
    # z_accel_offset = int(-755.5)
    # x_gyro_offset = int(16.5)
    # y_gyro_offset = int(-13.25)
    # z_gyro_offset = int(-25)
    
    
    if x_accel_offset != 0:
        set_x_accel_offset(x_accel_offset)
    
    if y_accel_offset != 0:
        set_y_accel_offset(y_accel_offset)
    
    if z_accel_offset != 0:
        set_z_accel_offset(z_accel_offset)
    
    if x_gyro_offset != 0:
        set_x_gyro_offset(x_gyro_offset)
    
    if y_gyro_offset != 0:
        set_y_gyro_offset(y_gyro_offset)
    
    if z_gyro_offset != 0:
        set_z_gyro_offset(z_gyro_offset)
    
    
    return gyro_config_vals[gyro_indx],accel_config_vals[accel_indx]
    
def read_raw_bits(register):
    # read accel and gyro values
    high = bus.read_byte_data(MPU6050_ADDR, register)
    low = bus.read_byte_data(MPU6050_ADDR, register+1)

    # combine higha and low for unsigned bit value
    value = ((high << 8) | low)
    
    # convert to +- value
    if(value > 32768):
        value -= 65536
    return value

def mpu6050_conv():
    # raw acceleration bits
    acc_x = read_raw_bits(ACCEL_XOUT_H)
    acc_y = read_raw_bits(ACCEL_YOUT_H)
    acc_z = read_raw_bits(ACCEL_ZOUT_H)

    # raw temp bits
##    t_val = read_raw_bits(TEMP_OUT_H) # uncomment to read temp
    
    # raw gyroscope bits
    gyro_x = read_raw_bits(GYRO_XOUT_H)
    gyro_y = read_raw_bits(GYRO_YOUT_H)
    gyro_z = read_raw_bits(GYRO_ZOUT_H)



    #convert to acceleration in g and gyro dps
    a_x = (acc_x/(2.0**15.0))*accel_sens
    a_y = (acc_y/(2.0**15.0))*accel_sens
    a_z = (acc_z/(2.0**15.0))*accel_sens

    w_x = (gyro_x/(2.0**15.0))*gyro_sens
    w_y = (gyro_y/(2.0**15.0))*gyro_sens
    w_z = (gyro_z/(2.0**15.0))*gyro_sens

##    temp = ((t_val)/333.87)+21.0 # uncomment and add below in return
    return a_x,a_y,a_z,w_x,w_y,w_z
    
def mpu6050_raw_data():
    # raw acceleration bits
    acc_x = read_raw_bits(ACCEL_XOUT_H)
    acc_y = read_raw_bits(ACCEL_YOUT_H)
    acc_z = read_raw_bits(ACCEL_ZOUT_H)

    # raw gyroscope bits
    gyro_x = read_raw_bits(GYRO_XOUT_H)
    gyro_y = read_raw_bits(GYRO_YOUT_H)
    gyro_z = read_raw_bits(GYRO_ZOUT_H)
    return acc_x,acc_y,acc_z,gyro_x,gyro_y,gyro_z
   
# 互补滤波的时间常数
alpha = 0.98

# 上一个时间点的姿态角，初始为0
prev_roll = 0.0
prev_pitch = 0.0
prev_yaw = 0.0
prev_time = time.time()
# 姿态角计算
def calculate_attitude(a_x, a_y, a_z, w_x, w_y, w_z):
    global prev_roll, prev_pitch, prev_yaw, prev_time

    # 获取当前时间
    curr_time = time.time()
    dt = curr_time - prev_time

    # 计算加速度计得出的滚转角和俯仰角（单位：弧度）
    accel_roll = math.atan2(a_y, a_z)
    accel_pitch = math.atan2(-a_x, math.sqrt(a_y**2 + a_z**2))

    # 用陀螺仪数据积分得到滚转角、俯仰角和航向角的变化量
    gyro_roll_rate = w_x * dt
    gyro_pitch_rate = w_y * dt
    gyro_yaw_rate = w_z * dt

    # 使用互补滤波器结合加速度计和陀螺仪数据
    roll = alpha * (prev_roll + gyro_roll_rate) + (1 - alpha) * accel_roll
    pitch = alpha * (prev_pitch + gyro_pitch_rate) + (1 - alpha) * accel_pitch

    # 航向角只能依赖陀螺仪积分计算，因为没有磁力计
    yaw = prev_yaw + gyro_yaw_rate

    # 更新上一个时间点的姿态角
    prev_roll = roll
    prev_pitch = pitch
    prev_yaw = yaw
    prev_time = curr_time

    # 将角度转换为度数
    roll_deg = math.degrees(roll)
    pitch_deg = math.degrees(pitch)
    yaw_deg = math.degrees(yaw)

    return roll_deg, pitch_deg, yaw_deg

# MPU6050 Registers
MPU6050_ADDR = 0x68
PWR_MGMT_1   = 0x6B
SMPLRT_DIV   = 0x19
CONFIG       = 0x1A
GYRO_CONFIG  = 0x1B
ACCEL_CONFIG = 0x1C
INT_ENABLE   = 0x38
ACCEL_XOUT_H = 0x3B
ACCEL_YOUT_H = 0x3D
ACCEL_ZOUT_H = 0x3F
TEMP_OUT_H   = 0x41
GYRO_XOUT_H  = 0x43
GYRO_YOUT_H  = 0x45
GYRO_ZOUT_H  = 0x47




# start I2C driver
bus = smbus2.SMBus(0) # start comm with i2c bus
gyro_sens,accel_sens = MPU6050_start() # instantiate gyro/accel


def mean_sensors():
    buffersize=1000
    buff_ax=0
    buff_ay=0
    buff_az=0
    buff_gx=0
    buff_gy=0
    buff_gz=0;
    i = 0
    while i<buffersize+101:
        try:
            raw = mpu6050_raw_data()
            a_raw = raw[0:3]
            g_raw = raw[3:6]
        except:
            continue


        if i>100 and i<=(buffersize+101):
            buff_ax=buff_ax+a_raw[0]
            buff_ay=buff_ay+a_raw[1]
            buff_az=buff_az+a_raw[2]
            buff_gx=buff_gx+g_raw[0]
            buff_gy=buff_gy+g_raw[1]
            buff_gz=buff_gz+g_raw[2]
        if (i==(buffersize+100)):
            mean_ax=buff_ax/buffersize
            mean_ay=buff_ay/buffersize
            mean_az=buff_az/buffersize
            mean_gx=buff_gx/buffersize
            mean_gy=buff_gy/buffersize
            mean_gz=buff_gz/buffersize

        i = i+1
        time.sleep(0.002)
    return int(mean_ax),int(mean_ay),int(mean_az),int(mean_gx),int(mean_gy),int(mean_gz)


def calibration(means_out):
    acel_deadzone=8
    giro_deadzone=1
    mean_ax = means_out[0]
    mean_ay = means_out[1]
    mean_az = means_out[2]
    mean_gx = means_out[3]
    mean_gy = means_out[4]
    mean_gz = means_out[5]

    ax_offset=-mean_ax/8
    ay_offset=-mean_ay/8
    az_offset=(16384-mean_az)/8
    gx_offset= -mean_gx/4
    gy_offset= -mean_gy/4
    gz_offset= -mean_gz/4

    while True:
        ready = 0
        set_x_accel_offset(int(ax_offset))
        set_y_accel_offset(int(ay_offset))
        set_z_accel_offset(int(az_offset))

        set_x_gyro_offset(int(gx_offset))
        set_y_gyro_offset(int(gy_offset))
        set_z_gyro_offset(int(gz_offset))

        means_out = mean_sensors()
        mean_ax = means_out[0]
        mean_ay = means_out[1]
        mean_az = means_out[2]
        mean_gx = means_out[3]
        mean_gy = means_out[4]
        mean_gz = means_out[5]

        print("...")
        print(means_out)
        print(ax_offset,ay_offset,az_offset,gx_offset,gy_offset,gz_offset)

        if abs(mean_ax)<= acel_deadzone:
            ready = ready+1
        else:
            ax_offset=ax_offset-mean_ax/acel_deadzone
    
        if abs(mean_ay)<= acel_deadzone:
            ready = ready+1
        else:
            ay_offset=ay_offset-mean_ay/acel_deadzone
    
        if abs(mean_az)<= acel_deadzone:
            ready = ready+1
        else:
            az_offset=az_offset-mean_az/acel_deadzone
    
        if abs(mean_gx)<= giro_deadzone:
            ready = ready+1
        else:
            gx_offset=gx_offset-mean_gx/(giro_deadzone+1)
    
        if abs(mean_gy)<= giro_deadzone:
            ready = ready+1
        else:
            gy_offset=gy_offset-mean_gy/(giro_deadzone+1)
    
        if abs(mean_gz)<= giro_deadzone:
            ready = ready+1
        else:
            gz_offset=gz_offset-mean_gz/(giro_deadzone+1)

        if ready==6:
            break
    return ax_offset,ay_offset,az_offset,gx_offset,gy_offset,gz_offset




if __name__ == "__main__":
    state =0

    # while True:
    #     if state ==0:
    #         print("Reading sensors for first time...")
    #         means_out = mean_sensors()
    #         state = state +1
    #         time.sleep(1)
    #     elif state == 1:
    #         print("\nCalculating offsets...")
    #         offsets = calibration(means_out)
    #         state = state + 1
    #         time.sleep(1)
    #     elif state == 2:
    #         means_out = mean_sensors()
    #         print("means out")
    #         print(means_out)
    #         print("offsets")
    #         print(offsets)
    #         break


    while True:

        # 读取传感器数据
        a_x, a_y, a_z, w_x, w_y, w_z = mpu6050_conv()
        # 计算姿态角
        # roll, pitch, yaw = calculate_attitude(a_x, a_y, a_z, w_x, w_y, w_z)

        str_show = "{0:.2f} {1:.2f} {2:.2f} {3:.2f} {4:.2f} {5:.2f}".format(a_x,a_y,a_z,w_x,w_y,w_z)

        print("\r %s"%(str_show),end='')
        # print(calculate_attitude(a_x, a_y, a_z, w_x, w_y, w_z))
        
        
